
import java.awt.*;
import java.awt.image.*;
import java.io.*;
import java.util.Collection;
import java.util.HashMap;
import java.util.HashSet;
import java.util.Iterator;
import java.util.Map;
import java.util.Set;
import java.util.SortedMap;
import java.util.TreeMap;
import javax.swing.*;

public class imageReader {
    static int colorBookSize = 64;
    static String fileName = "image4.rgb";//args[0];
    static int width = 352;//Integer.parseInt(args[5]);
    static int height = 288;//Integer.parseInt(args[6]);
    static byte[] bytes = null;
    static BufferedImage img = null;

    public static void main(String[] args) {
        fileName = args[0];
        colorBookSize = Integer.parseInt(args[1]);
        int method = Integer.parseInt(args[2]);

        img = new BufferedImage(width *2, height*2, BufferedImage.TYPE_INT_RGB);
        readFile();
        if(method == 1)
            UniformQuantization();
        if(method == 2)
            MedianCut();

        // Use a label to display the image
        JFrame frame = new JFrame();
        JLabel label = new JLabel(new ImageIcon(img));
        frame.getContentPane().add(label, BorderLayout.CENTER);
        frame.pack();
        frame.setVisible(true);

    }

    public static void readFile(){
        try {
            File file = new File(fileName);
            InputStream is = new FileInputStream(file);

            long len = file.length();
            bytes = new byte[(int) len];

            int offset = 0;
            int numRead = 0;
            while (offset < bytes.length && (numRead = is.read(bytes, offset, bytes.length - offset)) >= 0) {
                offset += numRead;
            }
            is.close();
        } catch (FileNotFoundException e) {
            e.printStackTrace();
            System.exit(0);
        } catch (IOException e) {
            e.printStackTrace();
            System.exit(0);
        }
    }

    public static void UniformQuantization(){
        int[][][] pixelCount = new int[64][64][64];
        int[][] pixelColors = new int[width][height];
        int ind = 0;
            for (int y = 0; y < height; y++) {
                for (int x = 0; x < width; x++) {
                    byte a = 0;
                    byte r = bytes[ind];
                    byte g = bytes[ind + height * width];
                    byte b = bytes[ind + height * width * 2];

                    int pix = 0xff000000 | ((r & 0xff) << 16) | ((g & 0xff) << 8) | (b & 0xff);
                    //int pix = ((a << 24) + (r << 16) + (g << 8) + b);
                    //img.setRGB(x, y, pix);
                    pixelColors[x][y] = pix;

                    Color c = new Color(pix);
                    int red = c.getRed()/4;
                    int blue = c.getBlue()/4;
                    int green = c.getGreen()/4;
                    pixelCount[red][green][blue]++;
                    ind++;
                }
            }

            SortedMap<Integer, Set> sortedRegions = new TreeMap();
            for(int r=0; r<64; r++)
                for(int g=0; g<64; g++)
                    for(int b=0; b<64; b++){
                        int count = pixelCount[r][g][b];
                        if(sortedRegions.containsKey(count)){
                            sortedRegions.get(count).add(new int[]{r,g,b});
                        }else{
                            Set region = new HashSet();
                            region.add(new int[]{r,g,b});
                            sortedRegions.put(count, region);
                        }
                    }

            int[][] colorBook = new int[colorBookSize][3];
            int colorSelected = 0;
            // regionsLeft = new HashSet();
            while(!sortedRegions.isEmpty()){
                int largestKey = sortedRegions.lastKey();
                Set regions = sortedRegions.get(largestKey);
                Iterator it = regions.iterator();
                while(it.hasNext()){
                    int[] region = (int[])it.next();
//                    region[0] = region[0]*4+2;
//                    region[1] = region[1]*4+2;
//                    region[2] = region[2]*4+2;
                    colorBook[colorSelected] = region;
                    it.remove();
                    if(++colorSelected == colorBookSize)
                        break;
                }
                if(colorSelected  == colorBookSize){
                    break;
                }
                sortedRegions.remove(largestKey);
            }

            int[][][] regionIndexs = new int[64][64][64];
            for(int i=0; i<colorBookSize; i++){
                int[] rgb = colorBook[i];
                regionIndexs[rgb[0]][rgb[1]][rgb[2]] = i;
            }
            Collection unselectedRegions = sortedRegions.values();
            Iterator it = unselectedRegions.iterator();
            while(it.hasNext()){
                Set regions = (Set)it.next();
                Iterator it2 = regions.iterator();
                while(it2.hasNext()){
                    int[] rgb = (int[])it2.next();
                    regionIndexs[rgb[0]][rgb[1]][rgb[2]] = findCloestRegion(rgb, colorBook);
                }
            }

            for (int y = 0; y < height; y++) {
                for (int x = 0; x < width; x++) {
                    Color c = new Color(pixelColors[x][y]);
                    int red = c.getRed()/4;
                    int blue = c.getBlue()/4;
                    int green = c.getGreen()/4;
                    int[] rgb = colorBook[regionIndexs[red][green][blue]];
                    Color c2 = new Color(rgb[0]*4+1, rgb[1]*4+1, rgb[2]*4+1, 255);
                    img.setRGB(x, y, c2.getRGB());
                    int redDiff = 255-Math.abs(c2.getRed() - c.getRed());
                    int greenDiff = 255-Math.abs(c2.getGreen() - c.getGreen());
                    int blueDiff = 255-Math.abs(c2.getBlue() - c.getBlue());
                    Color rdiff = new Color(redDiff,redDiff,redDiff);
                    Color gdiff = new Color(greenDiff,greenDiff,greenDiff);
                    Color bdiff = new Color(blueDiff,blueDiff,blueDiff);
                    img.setRGB(x+width, y, rdiff.getRGB());
                    img.setRGB(x, y+height, gdiff.getRGB());
                    img.setRGB(x+width, y+height, bdiff.getRGB());
                }
            }
    }

    public static int findCloestRegion(int[] rgb, int[][] colorBook){
        int minIndex = 0;
        int rDiff = colorBook[0][0]-rgb[0];
        int gDiff = colorBook[0][1]-rgb[1];
        int bDiff = colorBook[0][2]-rgb[2];
        double mindist = Math.sqrt(rDiff*rDiff + gDiff*gDiff+ bDiff*bDiff);

        for(int i=1; i<colorBook.length; i++){
            rDiff = colorBook[i][0]-rgb[0];
            gDiff = colorBook[i][1]-rgb[1];
            bDiff = colorBook[i][2]-rgb[2];
            double dist = Math.sqrt(rDiff*rDiff + gDiff*gDiff+ bDiff*bDiff);
            if(dist < mindist){
                minIndex = i;
                mindist = dist;
            }
        }
        return minIndex;
    }

    public static void MedianCut(){
        Color[][] pixelColors = new Color[width][height];
        Map<Integer, Set> ColorGroups = new HashMap(); //<color, coordinates>
        int ind = 0;
        Set allCoordinates = new HashSet();
        for (int y = 0; y < height; y++) {
            for (int x = 0; x < width; x++) {
                byte a = 0;
                byte r = bytes[ind];
                byte g = bytes[ind + height * width];
                byte b = bytes[ind + height * width * 2];

                int pix = 0xff000000 | ((r & 0xff) << 16) | ((g & 0xff) << 8) | (b & 0xff);
                //int pix = ((a << 24) + (r << 16) + (g << 8) + b);
                //img.setRGB(x, y, pix);

                pixelColors[x][y] = new Color(pix);
                if(ColorGroups.containsKey(pix)){
                    ColorGroups.get(pix).add(new int[]{x,y});
                }else{
                    Set pixel = new HashSet();
                    pixel.add(new int[]{x,y});
                    ColorGroups.put(pix, pixel);
                }
                allCoordinates.add(new int[]{x,y});
                ind++;
            }
        }

        SortedMap<Integer, Set> regions = new TreeMap(); //<size, coordinates>
        Set RegionGroup = new HashSet();
        RegionGroup.add(allCoordinates);
        regions.put(allCoordinates.size(), RegionGroup);
        int regionCount = 1;

        while(regionCount < colorBookSize){
            Iterator it = regions.get(regions.lastKey()).iterator();
            Set largestRegion = (Set)it.next();
            it.remove();
            if(regions.get(regions.lastKey()).isEmpty())
                regions.remove(regions.lastKey());
            Set[] newRegions = splitRegion(largestRegion, pixelColors, ColorGroups);
            if(regions.containsKey(newRegions[0].size())){
                regions.get(newRegions[0].size()).add(newRegions[0]);
            }else{
                Set newGroup = new HashSet();
                newGroup.add(newRegions[0]);
                regions.put(newRegions[0].size(), newGroup);
            }

            if(regions.containsKey(newRegions[1].size())){
                regions.get(newRegions[1].size()).add(newRegions[1]);
            }else{
                Set newGroup = new HashSet();
                newGroup.add(newRegions[1]);
                regions.put(newRegions[1].size(), newGroup);
            }
            regionCount++;
        }

        for(Map.Entry<Integer, Set> entry : regions.entrySet()){
            Set pixelGroups = entry.getValue();
            Iterator it = pixelGroups.iterator();
            while(it.hasNext()){
                Set pixelGroup = (Set)it.next();
                if(pixelGroup.size() == 0)
                    break;
                Iterator it2 = pixelGroup.iterator();
                int avgRed = 0;
                int avgGreen = 0;
                int avgBlue = 0;
                while(it2.hasNext()){
                    int[] xy = (int[])it2.next();
                    Color c = pixelColors[xy[0]][xy[1]];
                    avgRed += c.getRed();
                    avgGreen += c.getGreen();
                    avgBlue += c.getBlue();
                }
                avgRed = avgRed/pixelGroup.size();
                avgGreen = avgGreen/pixelGroup.size();
                avgBlue = avgBlue/pixelGroup.size();
                Color avgColor = new Color(avgRed, avgGreen, avgBlue);

                it2 = pixelGroup.iterator();
                while(it2.hasNext()){
                    int[] xy = (int[])it2.next();
                    Color c = pixelColors[xy[0]][xy[1]];
                    img.setRGB(xy[0], xy[1], avgColor.getRGB());

                    int redDiff = 255-Math.abs(c.getRed() - avgRed);
                    int greenDiff = 255-Math.abs(c.getGreen() - avgGreen);
                    int blueDiff = 255-Math.abs(c.getBlue() - avgBlue);
                    Color rdiff = new Color(redDiff,redDiff,redDiff);
                    Color gdiff = new Color(greenDiff,greenDiff,greenDiff);
                    Color bdiff = new Color(blueDiff,blueDiff,blueDiff);
                    img.setRGB(xy[0]+width, xy[1], rdiff.getRGB());
                    img.setRGB(xy[0], xy[1]+height, gdiff.getRGB());
                    img.setRGB(xy[0]+width, xy[1]+height, bdiff.getRGB());
                }
            }

        }

    }

    public static Set[] splitRegion(Set pixelSet, Color[][] pixelColors, Map<Integer, Set> ColorGroups){
        Set[] splitedRegions = new HashSet[2];
        int maxRed, maxGreen, maxBlue;
        int minRed, minGreen, minBlue;
        Iterator it = pixelSet.iterator();
        int[] xy = (int[])it.next();
        Color c = pixelColors[xy[0]][xy[1]];
        maxRed = minRed = c.getRed();
        maxGreen = minGreen = c.getGreen();
        maxBlue = minBlue = c.getBlue();
        Map<Integer, Set> colorInRegion = new HashMap();
        while(it.hasNext()){
            xy = (int[])it.next();
            c = pixelColors[xy[0]][xy[1]];
            int red = c.getRed();
            if(red>maxRed) maxRed = red;
            if(red<minRed) minRed = red;

            int green = c.getGreen();
            if(green>maxGreen) maxGreen = green;
            if(green<minGreen) minGreen = green;

            int blue = c.getBlue();
            if(blue>maxBlue) maxBlue = blue;
            if(blue<minBlue) minBlue = blue;

            int rgb = c.getRGB();
            if(colorInRegion.containsKey(rgb)){
                colorInRegion.get(rgb).add(xy);
            }else{
                Set coordinate = new HashSet();
                coordinate.add(xy);
                colorInRegion.put(rgb, coordinate);
            }
        }

        int maxAxis = 0;
        int maxRange = maxRed - minRed;
        if((maxGreen-minGreen) > maxRange){
            maxAxis = 1;
            maxRange = maxGreen-minGreen;
        }
        if((maxBlue-minBlue) > maxRange){
            maxAxis = 2;
        }

        SortedMap<Integer, Set> sortedRegion = new TreeMap();
        it = pixelSet.iterator();
        while(it.hasNext()){
            xy = (int[])it.next();
            c = pixelColors[xy[0]][xy[1]];
            int[] rgb2 = new int[]{c.getRed(), c.getGreen(), c.getBlue(), c.getRGB()};

            if(sortedRegion.containsKey(rgb2[maxAxis])){
                sortedRegion.get(rgb2[maxAxis]).add(rgb2[3]);
            }else{
                Set colors = new HashSet();
                colors.add(rgb2[3]);
                sortedRegion.put(rgb2[maxAxis], colors);
            }
        }

        splitedRegions[0] = new HashSet();
        splitedRegions[1] = new HashSet();
        int regionToAdd = 0;
        for(Map.Entry<Integer, Set> entry : sortedRegion.entrySet()){
            Set colors = entry.getValue();
            it = colors.iterator();
            while(it.hasNext()){
                int color = (Integer)it.next();
                Set pixelGroup = ColorGroups.get(color);
                if((splitedRegions[0].size()+pixelGroup.size()) > pixelSet.size()/2){
                    regionToAdd = 1;
                }
                splitedRegions[regionToAdd].addAll(pixelGroup);
            }
        }

        return splitedRegions;
    }
}
